#!/usr/bin python3
# -*- coding: utf-8 -*-
# @Time    : 19-9-8 下午4:45
# @Author  : Aries
# @Site    : 
# @File    : demo.py
# @Software: PyCharm
import cv2
import os
import mmcv
import copy
import numpy as np
from mmcv.runner import load_checkpoint
from mmdet.models import build_detector
import time
from mmdet.apis import init_detector, inference_detector, show_result
classname = ["bottle","other"]
# cfg = mmcv.Config.fromfile('configs/faster_rcnn_r50_fpn_1x.py')
# cfg.model.pretrained = None

config_file = '/home/ubuntu/Python-projects/mmdetection/tools/configs/bottle_cascade_rcnn_r50_fpn.py'
checkpoint_file = '/home/ubuntu/Python-projects/mmdetection/tools/work_dirs/bottle_cascade_rcnn_r50_fpn/latest.pth'

# 初始化模型
model = init_detector(config_file, checkpoint_file)
score_thr=0.8

# test a single image
files = open("/home/ubuntu/Python-projects/mmdetection/data/ImageSets/Main/test.txt")
for index, file in enumerate(files):
    t1=time.time()
    image_path = os.path.join("/home/ubuntu/Python-projects/mmdetection/data/checkbottle/JPEGImages/"+file.strip() + '.jpeg')
    image = mmcv.imread(image_path)
    image_save = copy.deepcopy(image)
    w, h, c = image.shape
    img = cv2.resize(image, dsize=(480, 480))
    w_rate = w / 480
    h_rate = h / 480
    result = inference_detector(model, img)
    t2=time.time()
    print("%d run time:"%index,t2-t1)
    for index, bboxes in enumerate(result):
        if len(bboxes)==0:
            continue
        if score_thr > 0:
            assert bboxes.shape[1] == 5
            scores = bboxes[:, -1]
            inds = scores > score_thr
            bboxes = bboxes[inds, :]
        if index==0:
            cls='bottle'
        else:
            cls = 'other'
        for bbox in bboxes:
            bbox_int = bbox.astype(np.int32)
            if bbox_int[1] > 2 and bbox_int[3] < 478:
                xmin = int(bbox_int[0] * h_rate)
                ymin = int(bbox_int[1] * w_rate)
                xmax = int(bbox_int[2] * h_rate)
                ymax = int(bbox_int[3] * w_rate)

                cv2.rectangle(image_save, (xmin, ymin), (xmax, ymax), (0, 255, 0), thickness=3)
                cv2.putText(image_save, cls, (int(xmin), int(ymin) - 2),
                            cv2.FONT_HERSHEY_COMPLEX, 1.5, (0, 2555, 0),thickness=3)
                if index==1:
                    bottle = image[ymin:ymax, xmin:xmax]  # 截取每个瓶子的图片，[y1:y2,x1:x2]

                    cv2.imwrite('./out_img/' + file.strip() + '.jpeg', bottle)
# img = cv2.resize(img,(400,400))
# print(result)
# image=show_result(img, result,classname)
# image=cv2.resize(image,(1280,720))
# cv2.imwrite("output.jpeg",image)